# -*- coding: utf-8 -*-
"""
MACD金叉筛选器 - 零轴上方或靠近零轴的金叉
筛选规则：
1. macd_dif > macd_dea（金叉）
2. macd_dif > 0 或从昨日 <=0 上穿至今日 >0（零轴上方或刚突破零轴）
3. macd_hist >= 0 或比昨日上升
4. 当日成交量 >= 0.5 * 20日平均成交量
"""
import logging
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from database import DatabaseManager
from config import DATABASE_CONFIG
import warnings

# 忽略pandas警告
warnings.filterwarnings('ignore', category=pd.errors.SettingWithCopyWarning)
warnings.filterwarnings('ignore', category=FutureWarning)

# 配置日志
# 获取项目根目录
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
LOG_DIR = os.path.join(PROJECT_ROOT, 'logs')
os.makedirs(LOG_DIR, exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler(os.path.join(LOG_DIR, 'macd_golden_cross_screener.log'), encoding='utf-8')
    ]
)
logger = logging.getLogger(__name__)


class MACDGoldenCrossScreener:
    """MACD金叉筛选器"""
    
    def __init__(self):
        self.db = DatabaseManager(DATABASE_CONFIG)
    
    def get_stock_daily_data(self, stock_code: str, days: int = 200) -> pd.DataFrame:
        """获取股票日线数据"""
        sql = """
        SELECT trade_date, open_price, high_price, low_price, close_price, volume, change_pct
        FROM new_stock_daily 
        WHERE stock_code = %s 
        ORDER BY trade_date DESC 
        LIMIT %s
        """
        try:
            results = self.db.execute_query(sql, (stock_code, days))
            if not results:
                return pd.DataFrame()
            
            df = pd.DataFrame(results)
            df['trade_date'] = pd.to_datetime(df['trade_date'])
            
            # 转换为float类型
            price_columns = ['open_price', 'high_price', 'low_price', 'close_price']
            for col in price_columns:
                df[col] = pd.to_numeric(df[col], errors='coerce')
            
            df['volume'] = pd.to_numeric(df['volume'], errors='coerce')
            df['change_pct'] = pd.to_numeric(df['change_pct'], errors='coerce')
            
            df = df.sort_values('trade_date').reset_index(drop=True)
            return df
        except Exception as e:
            logger.error(f"获取股票 {stock_code} 数据失败: {e}")
            return pd.DataFrame()
    
    def calculate_ma(self, df: pd.DataFrame, periods: list = [20, 60, 120]) -> pd.DataFrame:
        """计算移动平均线"""
        if df.empty:
            return df
        
        for period in periods:
            if len(df) >= period:
                df[f'ma{period}'] = df['close_price'].rolling(window=period).mean()
        
        return df
    
    def calculate_ema(self, df: pd.DataFrame, period: int) -> pd.Series:
        """计算指数移动平均线"""
        return df['close_price'].ewm(span=period, adjust=False).mean()
    
    def calculate_macd(self, df: pd.DataFrame, fast_period: int = 12, 
                      slow_period: int = 26, signal_period: int = 9) -> pd.DataFrame:
        """
        计算MACD指标
        DIF = EMA(12) - EMA(26)
        DEA = EMA(DIF, 9)
        HIST = (DIF - DEA) * 2
        """
        if df.empty or len(df) < slow_period:
            return df
        
        # 计算EMA
        ema_fast = self.calculate_ema(df, fast_period)
        ema_slow = self.calculate_ema(df, slow_period)
        
        # 计算DIF
        df['macd_dif'] = ema_fast - ema_slow
        
        # 计算DEA
        df['macd_dea'] = df['macd_dif'].ewm(span=signal_period, adjust=False).mean()
        
        # 计算HIST（柱状图）
        df['macd_hist'] = (df['macd_dif'] - df['macd_dea']) * 2
        
        return df
    
    def calculate_rsi(self, df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
        """计算RSI相对强弱指数"""
        if df.empty or len(df) < period:
            return df
        
        # 计算价格变化
        delta = df['close_price'].diff()
        
        # 分离上涨和下跌
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        
        # 计算平均涨幅和跌幅
        avg_gain = gain.rolling(window=period).mean()
        avg_loss = loss.rolling(window=period).mean()
        
        # 计算RS和RSI
        rs = avg_gain / avg_loss
        df['rsi'] = 100 - (100 / (1 + rs))
        
        return df
    
    def calculate_week_macd(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        计算周线MACD
        将日线数据转换为周线数据，然后计算MACD
        """
        if df.empty or len(df) < 26:
            df['week_macd_dif'] = np.nan
            df['week_macd_dea'] = np.nan
            return df
        
        try:
            # 设置trade_date为索引
            df_week = df.copy()
            df_week.set_index('trade_date', inplace=True)
            
            # 转换为周线数据（取每周最后一个交易日的数据）
            week_data = df_week['close_price'].resample('W').last()
            
            # 计算周线MACD
            ema_fast = week_data.ewm(span=12, adjust=False).mean()
            ema_slow = week_data.ewm(span=26, adjust=False).mean()
            week_dif = ema_fast - ema_slow
            week_dea = week_dif.ewm(span=9, adjust=False).mean()
            
            # 将周线MACD映射回日线数据
            # 使用forward fill将周线数据填充到对应的日线数据
            week_dif_reindexed = week_dif.reindex(df_week.index, method='ffill')
            week_dea_reindexed = week_dea.reindex(df_week.index, method='ffill')
            
            df['week_macd_dif'] = week_dif_reindexed.values
            df['week_macd_dea'] = week_dea_reindexed.values
        except Exception as e:
            logger.warning(f"计算周线MACD失败: {e}")
            df['week_macd_dif'] = np.nan
            df['week_macd_dea'] = np.nan
        
        return df
    
    def calculate_indicators(self, df: pd.DataFrame) -> pd.DataFrame:
        """计算所有技术指标"""
        if df.empty or len(df) < 130:
            return df
        
        # 计算移动平均线
        df = self.calculate_ma(df, periods=[20, 60, 120])
        
        # 计算MACD
        df = self.calculate_macd(df)
        
        # 计算RSI
        df = self.calculate_rsi(df)
        
        # 计算周线MACD
        df = self.calculate_week_macd(df)
        
        # 计算前一日收盘价
        df['prev_close'] = df['close_price'].shift(1)
        
        # 计算20日平均成交量
        df['avg_vol_20'] = df['volume'].rolling(window=20).mean()
        
        # 计算10天前的MA60和MA120
        df['ma60_10d_ago'] = df['ma60'].shift(10)
        df['ma120_10d_ago'] = df['ma120'].shift(10)
        
        # 计算60日最低价
        df['lowest_low_60'] = df['low_price'].rolling(window=60).min()
        
        # 计算60日最低MACD DIF值
        df['lowest_macd_dif_60'] = df['macd_dif'].rolling(window=60).min()
        
        # 计算前一日的MACD HIST（用于判断是否上升）
        df['prev_macd_hist'] = df['macd_hist'].shift(1)
        
        # 计算前一日的MACD DIF（用于判断零轴突破）
        df['prev_macd_dif'] = df['macd_dif'].shift(1)
        
        # 计算前一日的MACD DEA（用于判断金叉时间）
        df['prev_macd_dea'] = df['macd_dea'].shift(1)
        
        # 计算前两日的MACD DIF和DEA（用于判断金叉时间）
        df['prev2_macd_dif'] = df['macd_dif'].shift(2)
        df['prev2_macd_dea'] = df['macd_dea'].shift(2)
        
        return df
    
    def check_golden_cross_condition(self, row: pd.Series) -> bool:
        """
        检查是否满足MACD金叉条件
        0. 金叉必须在最近2个交易日内出现（新增）
        1. macd_dif > macd_dea（金叉）
        2. macd_dif > 0 或从昨日 <=0 上穿至今日 >0（零轴上方或刚突破零轴）
        3. macd_hist >= 0 或比昨日上升
        4. 当日成交量 >= 0.5 * 20日平均成交量
        """
        try:
            # 检查必要字段是否存在
            required_fields = ['macd_dif', 'macd_dea', 'macd_hist', 'prev_macd_hist', 
                             'prev_macd_dif', 'prev_macd_dea', 'prev2_macd_dif', 
                             'prev2_macd_dea', 'volume', 'avg_vol_20']
            if any(pd.isna(row[field]) for field in required_fields):
                return False
            
            # 条件0：金叉必须在最近2个交易日内出现（新增）
            # 当前日必须是金叉状态
            if row['macd_dif'] <= row['macd_dea']:
                return False
            
            # 检查金叉是否是新发生的（前1日或前2日时还未金叉）
            # 前1日金叉 或 前2日金叉，则说明金叉已经持续了至少1-2天，不符合条件
            prev1_is_golden = row['prev_macd_dif'] > row['prev_macd_dea']
            prev2_is_golden = row['prev2_macd_dif'] > row['prev2_macd_dea']
            
            # 金叉必须是新发生的：
            # 情况1：昨日未金叉，今日金叉（今日刚金叉）
            # 情况2：前日未金叉，昨日金叉（昨日刚金叉）
            is_fresh_golden_cross = (not prev1_is_golden) or (not prev2_is_golden)
            
            if not is_fresh_golden_cross:
                return False
            
            # 条件2：零轴上方或刚突破零轴
            condition2 = (row['macd_dif'] > 0) or \
                        (row['prev_macd_dif'] <= 0 and row['macd_dif'] > 0)
            if not condition2:
                return False
            
            # 条件3：HIST >= 0 或比昨日上升
            condition3 = (row['macd_hist'] >= 0) or \
                        (row['macd_hist'] > row['prev_macd_hist'])
            if not condition3:
                return False
            
            # 条件4：成交量放大
            if row['volume'] < 0.5 * row['avg_vol_20']:
                return False
            
            return True
            
        except Exception as e:
            logger.error(f"检查金叉条件失败: {e}")
            return False
    
    def screen_stock(self, stock_code: str, stock_name: str) -> dict:
        """筛选单只股票，带重试机制"""
        max_retries = 2
        for attempt in range(max_retries):
            try:
                # 获取日线数据
                df = self.get_stock_daily_data(stock_code, days=200)
                if df.empty or len(df) < 130:
                    return None
                
                # 计算技术指标
                df = self.calculate_indicators(df)
                
                # 获取最新一天的数据
                latest = df.iloc[-1]
                
                # 检查是否满足金叉条件
                if not self.check_golden_cross_condition(latest):
                    return None
                
                # 返回筛选结果
                result = {
                    'stock_code': stock_code,
                    'stock_name': stock_name,
                    'trade_date': latest['trade_date'].strftime('%Y-%m-%d'),
                    'close_price': round(float(latest['close_price']), 2),
                    'macd_dif': round(float(latest['macd_dif']), 4),
                    'macd_dea': round(float(latest['macd_dea']), 4),
                    'macd_hist': round(float(latest['macd_hist']), 4),
                    'volume': int(latest['volume']),
                    'avg_vol_20': round(float(latest['avg_vol_20']), 2),
                    'vol_ratio': round(float(latest['volume'] / latest['avg_vol_20']), 2),
                    'ma20': round(float(latest['ma20']), 2) if not pd.isna(latest['ma20']) else None,
                    'ma60': round(float(latest['ma60']), 2) if not pd.isna(latest['ma60']) else None,
                    'ma120': round(float(latest['ma120']), 2) if not pd.isna(latest['ma120']) else None,
                    'rsi': round(float(latest['rsi']), 2) if not pd.isna(latest['rsi']) else None,
                    'lowest_low_60': round(float(latest['lowest_low_60']), 2) if not pd.isna(latest['lowest_low_60']) else None,
                    'lowest_macd_dif_60': round(float(latest['lowest_macd_dif_60']), 4) if not pd.isna(latest['lowest_macd_dif_60']) else None,
                    'week_macd_dif': round(float(latest['week_macd_dif']), 4) if not pd.isna(latest['week_macd_dif']) else None,
                    'week_macd_dea': round(float(latest['week_macd_dea']), 4) if not pd.isna(latest['week_macd_dea']) else None,
                }
                
                return result
                
            except Exception as e:
                if attempt < max_retries - 1:
                    logger.warning(f"筛选股票 {stock_code} 失败，重试中 (尝试 {attempt + 1}/{max_retries}): {e}")
                    import time
                    time.sleep(0.5)  # 短暂延迟后重试
                else:
                    logger.error(f"筛选股票 {stock_code} 失败: {e}")
                    return None
    
    def get_all_active_stocks(self):
        """获取所有活跃股票"""
        sql = """
        SELECT stock_code, stock_name 
        FROM new_stock_info 
        WHERE status = 'active'
        ORDER BY stock_code
        """
        return self.db.execute_query(sql)
    
    def run_screening(self, limit: int = None):
        """运行筛选"""
        logger.info("=" * 80)
        logger.info("开始MACD金叉筛选 - 零轴上方或靠近零轴的金叉")
        logger.info("=" * 80)
        
        # 获取所有活跃股票
        stocks = self.get_all_active_stocks()
        total_stocks = len(stocks)
        
        if limit:
            stocks = stocks[:limit]
            logger.info(f"总共 {total_stocks} 只股票，本次筛选前 {limit} 只")
        else:
            logger.info(f"总共 {total_stocks} 只股票")
        
        results = []
        success_count = 0
        error_count = 0
        
        for idx, stock in enumerate(stocks, 1):
            stock_code = stock['stock_code']
            stock_name = stock['stock_name']
            
            if idx % 100 == 0:
                logger.info(f"进度: {idx}/{len(stocks)} ({idx*100//len(stocks)}%)")
            
            try:
                result = self.screen_stock(stock_code, stock_name)
                if result:
                    results.append(result)
                    logger.info(f"✓ 找到符合条件的股票: {stock_code} {stock_name}")
                success_count += 1
                
                # 添加小延迟，避免过于频繁的数据库查询
                if idx % 50 == 0:
                    import time
                    time.sleep(0.1)
                    
            except Exception as e:
                error_count += 1
                logger.error(f"✗ 处理股票 {stock_code} 失败: {e}")
        
        logger.info("=" * 80)
        logger.info(f"筛选完成！")
        logger.info(f"处理成功: {success_count} 只")
        logger.info(f"处理失败: {error_count} 只")
        logger.info(f"符合条件: {len(results)} 只")
        logger.info("=" * 80)
        
        # 保存结果
        if results:
            self.save_results(results)
        else:
            logger.info("没有找到符合条件的股票")
        
        return results
    
    def save_results(self, results: list):
        """保存筛选结果到CSV文件"""
        if not results:
            return
        
        # 生成文件名
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        csv_filename = f'macd_golden_cross_{timestamp}.csv'
        txt_filename = f'stock_codes_{timestamp}.txt'
        
        # 保存CSV文件 - 只保存股票代码，用逗号分隔
        stock_codes = [result['stock_code'] for result in results]
        with open(csv_filename, 'w', encoding='utf-8-sig') as f:
            f.write(','.join(stock_codes))
        logger.info(f"结果已保存到: {csv_filename}")
        
        # 保存股票代码到txt文件
        with open(txt_filename, 'w', encoding='utf-8') as f:
            for result in results:
                f.write(f"{result['stock_code']}\n")
        logger.info(f"股票代码已保存到: {txt_filename}")
        
        # 打印前10条结果
        logger.info("\n前10条筛选结果:")
        logger.info("-" * 80)
        for i, result in enumerate(results[:10], 1):
            logger.info(f"{i}. {result['stock_code']} {result['stock_name']} "
                       f"收盘价:{result['close_price']} "
                       f"DIF:{result['macd_dif']:.4f} "
                       f"DEA:{result['macd_dea']:.4f} "
                       f"量比:{result['vol_ratio']:.2f}")


def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description='MACD金叉筛选器')
    parser.add_argument('--limit', type=int, help='限制筛选的股票数量（用于测试）')
    parser.add_argument('--stock', type=str, help='筛选指定股票代码')
    
    args = parser.parse_args()
    
    screener = MACDGoldenCrossScreener()
    
    if args.stock:
        # 筛选单只股票
        logger.info(f"筛选单只股票: {args.stock}")
        sql = "SELECT stock_name FROM new_stock_info WHERE stock_code = %s"
        stock_info = screener.db.execute_query(sql, (args.stock,))
        if stock_info:
            stock_name = stock_info[0]['stock_name']
            result = screener.screen_stock(args.stock, stock_name)
            if result:
                logger.info(f"符合条件: {result}")
            else:
                logger.info("不符合筛选条件")
        else:
            logger.error(f"未找到股票: {args.stock}")
    else:
        # 批量筛选
        screener.run_screening(limit=args.limit)


if __name__ == '__main__':
    main()

